# -*- coding: utf-8 -*-
# time: 2025/3/27 16:56
# file: chromadb_ollama_ARG.py
# author: hanson
"""
https://juejin.cn/post/7424431087894380584
先参考 ch05_falsk_qwen.py 文件
Ollama支持embedding models嵌入模型
ollama+mxbai-embed-large把documents生成向量嵌入并写入向量库chromadb
"""
import chromadb.config
import ollama
documents = [
    "咖啡树是一种茜草科的常绿灌木或小乔木，原产于非洲的埃塞俄比亚和苏丹",
    "咖啡树的果实被称为咖啡樱桃，成熟时会变成鲜红色或紫色，每个果实通常包含两颗咖啡豆",
    "咖啡豆是咖啡树果实内的种子，经过烘焙后可以研磨成粉末冲泡咖啡",
    "阿拉比卡咖啡和罗布斯塔咖啡是两种主要的咖啡品种，阿拉比卡咖啡口感更佳，而罗布斯塔咖啡咖啡因含量更高",
    "咖啡种植需要温暖的气候和充足的降雨，咖啡树通常生长在海拔1000米到2000米的山区",
    "咖啡是世界上最受欢迎的饮料之一，具有提神醒脑的作用，并且富含抗氧化剂",
]

client = chromadb.Client()
collection = client.create_collection(name="docs")

# 将每个文档存储在向量嵌入数据库中
for i, d in enumerate(documents):
    response = ollama.embeddings(model="mxbai-embed-large", prompt=d)
    embedding = response["embedding"]
    collection.add(ids=[str(i)], embeddings=[embedding], documents=[d])
# 一个示例提示
prompt = "咖啡树的果实被称为什么？"

# 为提示生成一个嵌入并检索最相关的文档。
response = ollama.embeddings(prompt=prompt, model="mxbai-embed-large")
results = collection.query(query_embeddings=[response["embedding"]], n_results=3)
print(results)
